import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
import random
from config import BATCH_SIZE, LR, GAMMA, EPS_START, EPS_END, EPS_DECAY, MEMORY_CAPACITY, TARGET_UPDATE


class DQN(nn.Module):
    """DQN神经网络"""

    def __init__(self, input_dim, output_dim):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


class DQNAgent:
    def __init__(self, state_dim, action_dim):
        self.state_dim = state_dim  # 状态维度
        self.action_dim = action_dim  # 动作维度（每个资产3个动作，总动作数=3*n_assets）

        # 构建策略网络和目标网络
        self.policy_net = DQN(state_dim, action_dim)
        self.target_net = DQN(state_dim, action_dim)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        # 优化器和损失函数
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=LR)
        self.loss_fn = nn.MSELoss()

        # 经验回放池
        self.memory = deque(maxlen=MEMORY_CAPACITY)

        # 探索率相关
        self.eps = EPS_START
        self.eps_decay = EPS_DECAY

    def select_action(self, state):
        """选择动作：epsilon-greedy策略"""
        if random.random() > self.eps:
            # 贪心动作（基于策略网络）
            # 先将state转为numpy数组再转张量（避免单个数组转张量的潜在问题）
            state = torch.tensor(np.array(state), dtype=torch.float32).unsqueeze(0)
            with torch.no_grad():
                q_values = self.policy_net(state)
            action = q_values.argmax(dim=1).numpy()[0]
        else:
            # 随机动作
            action = random.randint(0, self.action_dim - 1)

        # 衰减探索率（注意：这里原逻辑是线性衰减，若需要指数衰减可调整）
        self.eps = max(EPS_END, self.eps - (EPS_START - EPS_END) / self.eps_decay)
        return action

    def store_transition(self, state, action, reward, next_state, done):
        """存储经验到回放池"""
        self.memory.append((state, action, reward, next_state, done))

    def learn(self):
        """从经验中学习（核心修改部分）"""
        if len(self.memory) < BATCH_SIZE:
            return  # 回放池数据不足，不学习

        # 采样批次数据
        batch = random.sample(self.memory, BATCH_SIZE)
        state_batch, action_batch, reward_batch, next_state_batch, done_batch = zip(*batch)

        # 关键修改：先将列表转为numpy数组，再转张量（消除警告）
        state_batch = torch.tensor(np.array(state_batch), dtype=torch.float32)
        action_batch = torch.tensor(np.array(action_batch), dtype=torch.long).unsqueeze(1)  # 保持unsqueeze(1)
        reward_batch = torch.tensor(np.array(reward_batch), dtype=torch.float32)
        next_state_batch = torch.tensor(np.array(next_state_batch), dtype=torch.float32)
        done_batch = torch.tensor(np.array(done_batch), dtype=torch.float32)

        # 计算当前Q值（策略网络）
        current_q = self.policy_net(state_batch).gather(1, action_batch).squeeze(1)

        # 计算目标Q值（目标网络）
        with torch.no_grad():
            next_q = self.target_net(next_state_batch).max(1)[0]
            target_q = reward_batch + GAMMA * next_q * (1 - done_batch)  # done时不累积未来奖励

        # 计算损失并优化
        loss = self.loss_fn(current_q, target_q)
        self.optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播
        self.optimizer.step()  # 更新参数

        return loss.item()

    def update_target_net(self):
        """更新目标网络参数"""
        self.target_net.load_state_dict(self.policy_net.state_dict())